package org.heapbd.simulation;

import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;

/**
 * This class generates IP netflow records by simulating a large area network.
 * 
 * It generates the desired number of records, using multiple threads for 
 * improved performance.
 * 
 * @author Levi
 *
 */
public class Simulator {
	
	public static void printUsage()
	{
		System.err.println("USAGE: Simulator <outputDir> " +
				"<outputRecordCount> <threadCount> <gzipOutput>");
		System.err.println("where <outputDir> is the location for files to be written to, ");
		System.err.println("outputRecordCount is the total number of netflow records generated, ");
		System.err.println("and threadCount is the number of parallel thread to use.  Each thread ");
		System.err.println("will produce its own output file in the outputDir.");
		System.err.println("If the gzipOutput parameter is set to true, the output files will");
		System.err.println("be gzipped.");
	}
	
	public static void main(String[] args)
	{
		if (args.length < 3)
		{
			printUsage();
			System.exit(-1);
		}
		
		Date start = new Date();
		System.out.println("Starting at " + start);
		
		generateData(args);
		
		Date end = new Date();
		System.out.println("Finished at " + end);
		System.out.println("Took " + ((end.getTime() - start.getTime()) / 1000) + " seconds.");
	}
	
	public static void generateData(String[] args)
	{
		String outputDir = args[0];
		
		File outputDirFile = new File(outputDir);
		if (! outputDirFile.exists())
		{
			outputDirFile.mkdirs();
		}
		
		String outputRecordCountString = args[1];
		String threadCountString = args[2];
		boolean gzipOutput = false;
		if (args.length > 3 && "true".equals(args[3]))
		{
			gzipOutput = true;
		}
		
		long outputRecordCount = 1000000000L;
		try {
			outputRecordCount = Long.valueOf(outputRecordCountString);
		} catch (NumberFormatException nfe) { 
			System.err.println("Couldn't interpret " + outputRecordCountString 
					+ " as a long, using " + outputRecordCount);
		}
		
		int threadCount = 1;
		try {
			threadCount = Integer.valueOf(threadCountString);
		} catch (NumberFormatException nfe) {
			System.err.println("Couldn't interpret " + threadCountString 
					+ " as an int, using " + threadCount);
		}
		
		// generate set of client and server IP addresses, each with an activity score
		// client IP count will change depending on the amount of data to be generated
		long clientCount = outputRecordCount / 1000;
		
		Map<String, Integer> clientIps = getIps(clientCount, 123L);
		String[] clientArray = clientIps.keySet().toArray(new String[clientIps.size()]);
		Arrays.sort(clientArray);
		Map<String, Integer> serverIps = getIps(10000, 234L);
		
		// sanity check IPs
		sanityCheck(clientIps, serverIps);
		
		// now generate client and server port profiles
		Integer[] clientPorts = getPorts(49152, 65535, 345L);
		Integer[] serverPorts = getPorts(0, 1023, 456L);
		
		int maxSecondsAfter = (int) outputRecordCount / 10000;
		
		// now we can generate our netflow data, using the number of threads 
		// we've been assigned.  To do this, we'll split the client IP space
		// evenly between threads.  Each thread will get a copy of the 
		// server IPs and port profiles.
		
		Set<Thread> dataGenerationThreads = new HashSet<Thread>();
		int clientsPerThread = clientArray.length / threadCount;
		long outputRecordsPerThread = outputRecordCount / threadCount;
		
		for (int x = 0; x < threadCount; x++)
		{
			// get subset of clients assigned to this thread
			Map<String, Integer> threadClients = new HashMap<String, Integer>();
			for (int y = 0; y < clientsPerThread; y++)
			{
				String clientIp = clientArray[(x * clientsPerThread) + y];
				threadClients.put(clientIp, clientIps.get(clientIp));
			}
			
			DataGenerationThread dataGenerator = 
					new DataGenerationThread(threadClients, serverIps, clientPorts, serverPorts, 
							outputDir, x, outputRecordsPerThread, maxSecondsAfter, gzipOutput);
			Thread thread = new Thread(dataGenerator);
			thread.start();
			dataGenerationThreads.add(thread);
		}
		
		for (Thread thread : dataGenerationThreads)
		{
			try {
				thread.join();
			} catch (InterruptedException ie) { }
		}
	}
	
	public static Integer[] getPorts(int minPort, int maxPort, long randomSeed)
	{
		Random portRandomiser = new Random(randomSeed);
		
		List<Integer> ports = new ArrayList<Integer>();
		
		for (int x = minPort; x < maxPort; x++)
		{
			int portActivity = portRandomiser.nextInt(100);
			if (portRandomiser.nextInt(100) == 0) // one in 100 ports will be boosted 100-fold
			{
				portActivity *= 100;
			}
			for (int y = 0; y < portActivity; y++)
			{
				ports.add(x);
			}
		}
		
		return ports.toArray(new Integer[ports.size()]);
	}
	
	public static void sanityCheck(Map<String, Integer> clientIps, Map<String, Integer> serverIps)
	{
		System.out.println("Clients first:");
		String[] clients = clientIps.keySet().toArray(new String[clientIps.size()]);
		for (int x = 0; x < 100; x++)
		{
			System.out.println(clients[x] + " = " + clientIps.get(clients[x]));
		}
		System.out.println("Now servers:");
		String[] servers = serverIps.keySet().toArray(new String[serverIps.size()]);
		for (int x = 0; x < 100; x++)
		{
			System.out.println(servers[x] + " = " + serverIps.get(servers[x]));
		}
		
		for (String server : servers)
		{
			if (clientIps.containsKey(server))
			{
				System.out.println(server + " is client as well as server.");
			}
		}
	}
	
	/**
	 * Get a number of IP addresses using the supplied random seed.
	 * 
	 * The returned IP addresses are assigned an activity level, which
	 * can be used to determine their connection count, or their frequency
	 * of connection establishment, or both.
	 * 
	 * @param ipCount
	 * @param randomSeed
	 * @return
	 */
	public static Map<String, Integer> getIps(long ipCount, long randomSeed)
	{
		Map<String, Integer> result = new HashMap<String, Integer>();
		
		Random random = new Random(randomSeed);
		
		// generate IPs
		for (int x = 0; x < ipCount; x++)
		{
			String ipAddr = getIp(random);
			while (result.containsKey(ipAddr))
			{
				ipAddr = getIp(random);
			}
			
			// activity level is between 1 and 1000
			Integer activity = random.nextInt(1000) + 1;
			if (random.nextInt(1000) == 0)
			{
				activity *= 1000; // one in 1000 will get boosted by 1000
			}
			result.put(ipAddr, activity);
		}
		
		return result;
	}
	
	/**
	 * Generate an IP address from the supplied random number generator.
	 * Note that no distinction is made between private or public IP space.
	 * 
	 * @param random
	 * @return
	 */
	public static String getIp(Random random)
	{
		int firstOctet = random.nextInt(254) + 1;
		int secondOctet = random.nextInt(254) + 1;
		int thirdOctet = random.nextInt(254) + 1;
		int fourthOctet = random.nextInt(254) + 1;
		
		return firstOctet + "." + secondOctet + "." + thirdOctet + "." + fourthOctet;
	}
}
